import sys
sys.path.append("/home/bachelor/zhangyi/my_code/vasp")

import os
import subprocess
import time 
from utils import *
from datetime import datetime
import argparse
import logging

logging.basicConfig(level=logging.DEBUG, #设置日志输出格式
                    filename=f"/home/bachelor/zhangyi/my_code/vasp/log/{datetime.now()}.log", #log日志输出的文件位置和文件名
                    filemode="w", #文件的写入格式，w为重新写入文件，默认是追加
                    format="%(asctime)s - %(levelname)-9s - %(filename)-8s : %(lineno)s line - %(message)s", #日志输出的格式
                    datefmt="%Y-%m-%d %H:%M:%S")

processNum = 4
parser = argparse.ArgumentParser()
parser.add_argument('--vaspin', type=str)

def commitBatchTask(rootPath):
    """
    批量提交VASP任务,每个任务串行

    Args:
        rootPath (str): 根目录
    """
    logging.debug(f"DEBUG:batch vasp task start, time:{datetime.now()}")
    # 排序并记录任务顺序
    taskList = sorted(getAllFolderAbsPathInFolder(rootPath))
    logging.debug(taskList)
    count = len(taskList)
    # 逐个提交任务
    for i, task in enumerate(taskList):
        logging.debug(f"DEBUG:try to start task is {task} ({i + 1}/{count}), task {i} is running")
        commitSingleTask(task)
    # 等待最后一个任务完成
    while not checkTaskFinish():
        time.sleep(60)
    logging.debug(f"DEBUG:batch vasp task done, time:{datetime.now()}")

def commitSingleTask(path):
    """
    提交一个vasp作业

    path: 目录下包含四个输出文件
    """
    os.chdir(path)
    # 检查当前任务是否已经计算过
    if os.path.exists(f"{path}/DOSCAR") and os.path.exists(f"{path}/PROCAR"):
        logging.debug(f"DEBUG:current task {path} is running or done")
        return
    # 检查上个任务是否计算完成,如果没有就休眠一段时间再检测
    while not checkTaskFinish():
        time.sleep(60)
    logging.debug(f"DEBUG:current task path:{path}, start time:{datetime.now()}")
    # 执行计算
    os.system(f"mpirun -np {processNum} vasp_std")

def checkTaskFinish():
    """
    检查当前是否有VASP任务在执行
    """
    command = "ps -u | grep vasp_std | wc -l"
    result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out = result.stdout.decode().replace("\n", "")
    # 搜索结果不足指定进程数,说明任务结束
    return int(out) < processNum


if __name__ == '__main__':
    args = parser.parse_args()
    if args.vaspin is None or args.vaspin == "":
        print("需要指定VASP输入文件的路径,如:python DoCalculate.py --vaspin /home/zhangyi/testPredict/noBandCof/cifs")
    else:
        os.system("ulimit -s unlimited")
        commitBatchTask(args.vaspin)